#include "stm32f4xx_hal.h"
#include <stdint.h>
#include "ff.h" // FatFs header file
#include "bootloader.h"
#include "usb_host.h"
#define U_DETECT_TIMEOUT 5000 // 5 seconds timeout for USB detection
#define USB_DRIVE_PATH "0:"
#define APP_BIN_PATH "0:/app.bin"

static uint8_t udisk_update_state = 0; 
static uint32_t start_time = 0;
static FILINFO fno;
static FIL firmware_file;
static FATFS UsbFS;
FRESULT res;
extern ApplicationTypeDef Appli_state;
typedef  void (*pFunction)(void);
pFunction Jump_To_Application;
uint32_t JumpAddress;

typedef enum {
    UDISK_UPDATE_IDLE = 0,
    UDISK_UPDATE_CHECK,
    UDISK_UPDATE_START,
    UDISK_UPDATE_IN_PROGRESS,
    UDISK_UPDATE_COMPLETE
} udisk_update_state_t;

typedef struct {
    uint32_t sector_addr;   // 扇区起始地址
    uint8_t  is_erased;     // 标记是否已擦除（1=已擦除，0=未擦除）
} FlashSectorState;

static FlashSectorState current_sector = {0};

HAL_StatusTypeDef iap_write_appbin(uint32_t flash_addr, uint8_t *file_buffer, uint32_t bytes_read);
void JumpToApp(uint32_t appAddr);

void udisk_update(void)
{

    uint8_t file_buffer[512]={0};
    uint32_t bytes_read;
    unsigned long total_bytes_read = 0;
    uint32_t flash_addr = APP_FLASH_ADDR;
    switch (udisk_update_state)
    {
        case UDISK_UPDATE_IDLE:
            udisk_update_state = UDISK_UPDATE_CHECK;
            start_time = HAL_GetTick();
            break;
        case UDISK_UPDATE_CHECK:
            // Check if USB is connected
            if (HAL_GetTick() - start_time > U_DETECT_TIMEOUT) {
                udisk_update_state = UDISK_UPDATE_COMPLETE;
                printf("USB detection timeout.\n");
            }
            else {
                if(Appli_state == APPLICATION_READY)
                {
                    res = f_mount(&UsbFS, USB_DRIVE_PATH, 1);
                    if (res == FR_OK) {
                        printf("USB detected and mounted successfully.\n");
                        res = f_stat(APP_BIN_PATH, &fno);
                        if(res == FR_OK) {
                            udisk_update_state = UDISK_UPDATE_START;
                            printf("Update file found: %s\n", fno.fname);
                        }
                        else
                        {
                            udisk_update_state = UDISK_UPDATE_COMPLETE;
                            printf("Update file not found.\n");
                        }
                    }
                }
            }
            break;
        case UDISK_UPDATE_START:
            res = f_open(&firmware_file, APP_BIN_PATH, FA_READ);
            if(res == FR_OK) {
                udisk_update_state = UDISK_UPDATE_IN_PROGRESS;
                printf("Starting firmware update from: %s\n", APP_BIN_PATH);
            }
            else
            {
                udisk_update_state = UDISK_UPDATE_COMPLETE;
                printf("Failed to open update file: %s\n", APP_BIN_PATH);
            }
            break;
        case UDISK_UPDATE_IN_PROGRESS:

            // Here you would typically read the file and perform the update
            while (total_bytes_read < fno.fsize) {
                // 读取数据块到缓冲区
                res = f_read(&firmware_file, file_buffer, 512, &bytes_read);
                if (res != FR_OK || bytes_read == 0) {
                    // 读取出错或到达文件末尾
                    printf("Error reading file or end of file reached.\n");
                    break;
                }
                // 写入到FLASH
                printf("Writing %lu bytes to flash at address 0x%08lX\n", bytes_read, flash_addr);
                res = iap_write_appbin(flash_addr, file_buffer, bytes_read);

                // 更新计数器和地址
                total_bytes_read += bytes_read;
                flash_addr += bytes_read;
                
                // 显示进度
                printf("Progress: %lu/%lu bytes (%.2f%%)\n", total_bytes_read, fno.fsize, (total_bytes_read * 100.0) / fno.fsize);
            }
            f_close(&firmware_file);
            udisk_update_state = UDISK_UPDATE_COMPLETE;
            printf("Firmware update completed successfully.\n");
            break;
        case UDISK_UPDATE_COMPLETE:
            udisk_update_state = UDISK_UPDATE_IDLE;
            printf("Rebooting to new application...\n");
            JumpToApp(APP_FLASH_ADDR);
            break;
        default:
            break;
    }
}


#define FLASH_USER_START_ADDR  APP_FLASH_ADDR  // 示例：从第128KB地址开始写入（根据实际应用调整）
#define FLASH_USER_SIZE 0xF0000
#define FLASH_WAITETIME  50000          //FLASH等待超时时间

//FLASH 扇区的起始地址
#define ADDR_FLASH_SECTOR_0     ((uint32_t)0x08000000) 	//扇区0起始地址, 16 Kbytes  
#define ADDR_FLASH_SECTOR_1     ((uint32_t)0x08004000) 	//扇区1起始地址, 16 Kbytes  
#define ADDR_FLASH_SECTOR_2     ((uint32_t)0x08008000) 	//扇区2起始地址, 16 Kbytes  
#define ADDR_FLASH_SECTOR_3     ((uint32_t)0x0800C000) 	//扇区3起始地址, 16 Kbytes  
#define ADDR_FLASH_SECTOR_4     ((uint32_t)0x08010000) 	//扇区4起始地址, 64 Kbytes  
#define ADDR_FLASH_SECTOR_5     ((uint32_t)0x08020000) 	//扇区5起始地址, 128 Kbytes  
#define ADDR_FLASH_SECTOR_6     ((uint32_t)0x08040000) 	//扇区6起始地址, 128 Kbytes  
#define ADDR_FLASH_SECTOR_7     ((uint32_t)0x08060000) 	//扇区7起始地址, 128 Kbytes  
#define ADDR_FLASH_SECTOR_8     ((uint32_t)0x08080000) 	//扇区8起始地址, 128 Kbytes  
#define ADDR_FLASH_SECTOR_9     ((uint32_t)0x080A0000) 	//扇区9起始地址, 128 Kbytes  
#define ADDR_FLASH_SECTOR_10    ((uint32_t)0x080C0000) 	//扇区10起始地址,128 Kbytes  
#define ADDR_FLASH_SECTOR_11    ((uint32_t)0x080E0000) 	//扇区11起始地址,128 Kbytes 


//获取某个地址所在的flash扇区
//addr:flash地址
//返回值:0~11,即addr所在的扇区
uint8_t STMFLASH_GetFlashSector(uint32_t addr)
{
	if(addr<ADDR_FLASH_SECTOR_1)return FLASH_SECTOR_0;
	else if(addr<ADDR_FLASH_SECTOR_2)return FLASH_SECTOR_1;
	else if(addr<ADDR_FLASH_SECTOR_3)return FLASH_SECTOR_2;
	else if(addr<ADDR_FLASH_SECTOR_4)return FLASH_SECTOR_3;
	else if(addr<ADDR_FLASH_SECTOR_5)return FLASH_SECTOR_4;
	else if(addr<ADDR_FLASH_SECTOR_6)return FLASH_SECTOR_5;
	else if(addr<ADDR_FLASH_SECTOR_7)return FLASH_SECTOR_6;
	else if(addr<ADDR_FLASH_SECTOR_8)return FLASH_SECTOR_7;
	else if(addr<ADDR_FLASH_SECTOR_9)return FLASH_SECTOR_8;
	else if(addr<ADDR_FLASH_SECTOR_10)return FLASH_SECTOR_9;
	else if(addr<ADDR_FLASH_SECTOR_11)return FLASH_SECTOR_10;   
	return FLASH_SECTOR_11;	
}


/**
 * @brief 将数据写入Flash
 * @param flash_addr  目标Flash地址（必须为扇区起始地址）
 * @param file_buffer 源数据缓冲区
 * @param bytes_read  数据长度（字节）
 * @retval HAL_OK 成功，其他值失败
 */
HAL_StatusTypeDef iap_write_appbin(uint32_t flash_addr, uint8_t *file_buffer, uint32_t bytes_read) {
    HAL_StatusTypeDef status;
    uint32_t SectorError = 0;
    FLASH_EraseInitTypeDef EraseInitStruct;
    uint32_t *pData = (uint32_t *)file_buffer;
    uint32_t address = flash_addr;
    uint32_t remaining = bytes_read;

    // 1. 检查地址和长度是否合法
    if (flash_addr < FLASH_USER_START_ADDR || 
        (flash_addr + bytes_read) > (FLASH_USER_START_ADDR + FLASH_USER_SIZE)) {
        return HAL_ERROR;
    }

    // 记录当前扇区
    if(current_sector.sector_addr != STMFLASH_GetFlashSector(flash_addr))
    {
        current_sector.sector_addr = STMFLASH_GetFlashSector(flash_addr);
        current_sector.is_erased = 0; // 清除擦除标志
    }
    // 2. 解锁Flash
    HAL_FLASH_Unlock();
	
    // 3. 擦除目标扇区（STM32F4需按扇区擦除）
    if(current_sector.is_erased == 0)
    {
        EraseInitStruct.TypeErase = FLASH_TYPEERASE_SECTORS; // 擦除类型: 扇区擦除
        EraseInitStruct.VoltageRange = FLASH_VOLTAGE_RANGE_3; // 3.3V供电
        EraseInitStruct.Sector = STMFLASH_GetFlashSector(flash_addr); // 根据flash_addr选择扇区
        EraseInitStruct.NbSectors = 1;                         // 擦除1个扇区
        status = HAL_FLASHEx_Erase(&EraseInitStruct, &SectorError);
        if (status != HAL_OK) {
            HAL_FLASH_Lock();
            return HAL_ERROR;
        }
        FLASH_WaitForLastOperation(FLASH_WAITETIME);            //等待上次操作完成
        current_sector.is_erased = 1;
    }

    // 4. 按32位（4字节）写入数据（STM32F4必须按32位写入）
    while (remaining >= 4) {
        status = HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, address, *(uint32_t *)pData);
        if (status != HAL_OK) {
			continue;
        }
        address += 4;
        pData += 1;
        remaining -= 4;
    }

    // 5. 处理剩余不足4字节的数据（按字节写入）
    if (remaining > 0) {
        uint8_t *pRemaining = (uint8_t *)pData;
        while (remaining--) {
            status = HAL_FLASH_Program(FLASH_TYPEPROGRAM_BYTE, address, *pRemaining);
            if (status != HAL_OK) {
				continue;
            }
            address++;
            pRemaining++;
        }
    }

    // 6. 校验写入的数据
    if (memcmp((void *)flash_addr, file_buffer, bytes_read) != 0) {
        HAL_FLASH_Lock();
        return HAL_ERROR;
    }

    // 7. 重新锁定Flash
    HAL_FLASH_Lock();
    return HAL_OK;
}



void JumpToApp(uint32_t appAddr) {
    // 1. 关闭所有中断
    // __disable_irq(); // 这一句关闭, APP记得打开
    

    // 2. 复位外设（关键步骤！）
    HAL_RCC_DeInit();           // 复位时钟配置（可选，APP会重新初始化）
    HAL_DeInit();               // 复位所有HAL外设

    // 3. 关闭特定外设（按需添加）
    // HAL_DMA_DeInit(&hdma_usart1_tx);  // 关闭DMA
    // HAL_UART_DeInit(&huart1);         // 关闭UART
    // SysTick->CTRL = 0;                // 禁用SysTick

    // 4. 设置堆栈指针并跳转
    uint32_t *pStack = (uint32_t *)appAddr;
    uint32_t *pReset = (uint32_t *)(appAddr + 4);
    __set_MSP(*pStack);               // 设置主堆栈指针
    ((void (*)(void))(*pReset))();    // 跳转到Reset_Handler
}























